import os
import argparse
import subprocess
import logging

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("run_simple.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def run_command(command):
    """
    运行命令并记录日志
    """
    logger.info(f"执行命令: {command}")
    process = subprocess.Popen(
        command,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True
    )
    stdout, stderr = process.communicate()
    
    if process.returncode != 0:
        logger.error(f"命令执行失败: {stderr}")
        raise Exception(f"命令执行失败: {stderr}")
    
    logger.info(f"命令执行成功: {stdout}")
    return stdout

def main(args):
    # 1. 数据集分割
    if args.split_dataset:
        logger.info("步骤1: 数据集分割")
        run_command("python dataset_split.py")
    
    # 2. 文本分段处理
    if args.split_text:
        logger.info("步骤2: 文本分段处理")
        run_command(f"python split_text_data.py --filename train.csv --max_length {args.max_length}")
        run_command(f"python split_text_data.py --filename valid.csv --max_length {args.max_length}")
    
    # 3. 训练模型
    if args.train:
        logger.info("步骤3: 训练BertTextCNN模型")
        train_cmd = f"python train.py --train_file train.csv --valid_file valid.csv --bert_model_path {args.bert_model_path} --output_dir {args.output_dir} --max_length {args.max_length} --batch_size {args.batch_size} --epochs {args.epochs} --filter_sizes {args.filter_sizes} --num_filters {args.num_filters}"
        run_command(train_cmd)
    
    # 4. 模型预测
    if args.predict:
        logger.info("步骤4: 模型预测")
        predict_cmd = f"python predict.py --test_file {args.test_file} --output_file {args.output_file} --bert_model_path {args.bert_model_path} --model_path {args.model_path} --max_length {args.max_length} --batch_size 1 --filter_sizes {args.filter_sizes} --num_filters {args.num_filters}"
        run_command(predict_cmd)
    
    logger.info("所有步骤执行完成!")

if __name__ == "__main__":
    # 获取当前脚本所在目录的绝对路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 构建models文件夹的路径
    models_dir = os.path.join(current_dir, "models")
    
    parser = argparse.ArgumentParser(description="运行NLP新闻分类的简化流程")
    
    # 流程控制参数
    parser.add_argument("--split_dataset", action="store_true", help="是否执行数据集分割步骤")
    parser.add_argument("--split_text", action="store_true", help="是否执行文本分段处理步骤")
    parser.add_argument("--train", action="store_true", help="是否执行模型训练步骤")
    parser.add_argument("--predict", action="store_true", help="是否执行模型预测步骤")
    parser.add_argument("--all", action="store_true", help="执行所有步骤")
    
    # 通用参数
    parser.add_argument("--max_length", type=int, default=1024, help="最大序列长度")
    parser.add_argument("--batch_size", type=int, default=2, help="训练的批量大小")
    parser.add_argument("--epochs", type=int, default=3, help="训练轮数")
    parser.add_argument("--bert_model_path", type=str, default=models_dir, help="预训练BERT模型路径")
    parser.add_argument("--output_dir", type=str, default="./saved_models", help="模型保存目录")
    parser.add_argument("--filter_sizes", type=str, default="2,3,4", help="卷积核大小，用逗号分隔")
    parser.add_argument("--num_filters", type=int, default=50, help="每种卷积核的数量")
    
    # 预测参数
    parser.add_argument("--test_file", type=str, default="test_a.csv", help="测试数据文件路径")
    parser.add_argument("--output_file", type=str, default="prediction_result.csv", help="预测结果输出文件路径")
    parser.add_argument("--model_path", type=str, default="", help="BertTextCNN模型路径")
    
    args = parser.parse_args()
    
    # 如果指定了--all参数，则执行所有步骤
    if args.all:
        args.split_dataset = True
        args.split_text = True
        args.train = True
        args.predict = True
        
        # 设置默认的模型路径
        if not args.model_path:
            args.model_path = f"{args.output_dir}/best_model.pth"
    
    # 创建必要的目录
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    
    main(args) 